import cv2
import numpy as np


#调节图片
def adjust_brightness_contrast_saturation(image, brightness=0, contrast=0, saturation=0):
    # 调整亮度
    bright_image = cv2.convertScaleAbs(image, alpha=1, beta=brightness)

    # 调整对比度
    contrast_image = cv2.convertScaleAbs(bright_image, alpha=1 + contrast / 100, beta=0)

    # 调整饱和度
    hsv_image = cv2.cvtColor(contrast_image, cv2.COLOR_BGR2HSV)
    hsv_image[:, :, 1] = np.clip(hsv_image[:, :, 1] + saturation, 0, 255)  # 饱和度通道
    saturated_image = cv2.cvtColor(hsv_image, cv2.COLOR_HSV2BGR)

    return saturated_image


#轮廓识别
def detect_panels(image, color):
    if color == 'blue':
        lower_color = np.array([100, 150, 0])
        upper_color = np.array([140, 255, 255])
    elif color == 'red':
        lower_color = np.array([0, 150, 150])
        upper_color = np.array([10, 255, 255])
        lower_color2 = np.array([160, 150, 150])
        upper_color2 = np.array([180, 255, 255])

    hsv = cv2.cvtColor(image, cv2.COLOR_BGR2HSV)
  # 掩膜
    if color == 'red':
        mask1 = cv2.inRange(hsv,lower_color, upper_color)
        mask2 = cv2.inRange(hsv,lower_color2, upper_color2)
        mask = cv2.bitwise_or(mask1, mask2)
    else:
        mask = cv2.inRange(hsv,lower_color, upper_color)
  #去噪 形态学操作
    blurred = cv2.GaussianBlur(mask, (5, 5), 0)
    kernel = np.ones((5, 5), np.uint8)
    dilated = cv2.dilate(blurred, kernel, iterations=1)
    eroded = cv2.erode(dilated, kernel, iterations=1)
    cv2.imshow("s",eroded)
    contours, _ = cv2.findContours(eroded, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)
    #轮廓信息
    panels = []
    for contour in contours:
        if cv2.contourArea(contour) > 100:
            x, y, w, h = cv2.boundingRect(contour)
            panels.append((x, y, w, h))

    return panels


#判断两灯条是否平行
def are_panels_parallel(panel1, panel2):
    slope1 = panel1[3] / panel1[2]  # height / width
    slope2 = panel2[3] / panel2[2]  # height / width
    # 判断斜率差异（平行的条件是 |斜率1 - 斜率2| 应接近于 0）
    return abs(slope1 - slope2) < 5


#判断是否为装甲板灯条
def are_panel(panel1,panel2):
    # 计算长宽比
    aspect_ratio1 = panel1[2] / panel1[3]  # w/h
    aspect_ratio2 = panel2[2] / panel2[3]  # w/h

    # 检查长宽比是否在范围内（未测定）
    if not (1.5 <= aspect_ratio1 <= 1.6 or 1.5 <= aspect_ratio2 <= 1.6):
        return False


#调节图片大小
def resize_image(image, width=640):
    aspect_ratio = image.shape[1] / image.shape[0]
    height = int(width / aspect_ratio)
    return cv2.resize(image, (width, height))

def main():
    color = input("请输入要识别的灯条颜色 (blue/red): ").strip().lower()
    source_type = input("请输入输入源类型 (photo/video): ").strip().lower()
    adjust = input("是否需要调节图片(0:不需要，1：需要):").strip().lower()
    if adjust == "1":
        brightness = float(input("亮度:").strip())  # 亮度调整值
        contrast = float(input("对比度:").strip()) # 对比度调整值（百分比）
        saturation = float(input("饱和度:").strip())  # 饱和度调整值


    if source_type == 'photo':
        image_path = input("请输入照片路径: ")
        image = cv2.imread(image_path)
        image = resize_image(image)


        # 调整图像参数
        if adjust =="1":
         image = adjust_brightness_contrast_saturation(image, brightness, contrast, saturation)

        #提取灯条
        panels = detect_panels(image, color)
        if len(panels) >= 2:
            # 取前两个灯条
            panel1, panel2 = panels[0], panels[1]
            # 判断是否平行且为装甲板灯条长宽比
            if (are_panels_parallel(panel1, panel2)):
                # and are_panels(panel1,panel2))
                # 计算灯条中点
                midpoints = []
                for panel in [panel1, panel2]:
                    mid_x = panel[0] + panel[2] // 2
                    mid_y = panel[1] + panel[3] // 2
                    midpoints.append((mid_x, mid_y))
                 # 计算装甲板中心
                center_x = (midpoints[0][0] + midpoints[1][0]) // 2
                center_y = (midpoints[0][1] + midpoints[1][1]) // 2
                cv2.circle(image, (center_x,center_y), 20, (0, 255, 0), -1)  # 绘制实心绿色圆
            #绘制灯条中心
            for panel in panels:
                x, y, w, h = panel
                cv2.circle(image, (x + w // 2, y + h // 2),5, (255,255,0),-1)

        #结果展示
        cv2.imshow("Detected Panels", image)
        cv2.waitKey(0)
        cv2.destroyAllWindows()

    elif source_type == 'video':
        video_path = input("请输入视频路径: ")
        cap = cv2.VideoCapture(video_path)

        while cap.isOpened():
            ret, frame = cap.read()
            if not ret:
                break

            frame = resize_image(frame)
            if adjust == "1":
                frame = adjust_brightness_contrast_saturation(frame, brightness, contrast, saturation)
            panels = detect_panels(frame, color)

            if len(panels) >= 2:
                panel1, panel2 = panels[0], panels[1]
                if  (are_panels_parallel(panel1, panel2) and are_panel(panel1,panel2)):
                    midpoints = []
                    for panel in [panel1, panel2]:
                        mid_x = panel[0] + panel[2] // 2
                        mid_y = panel[1] + panel[3] // 2
                        midpoints.append((mid_x, mid_y))

                    center_x = (midpoints[0][0] + midpoints[1][0]) // 2
                    center_y = (midpoints[0][1] + midpoints[1][1]) // 2
                    cv2.circle(frame,(center_x,center_y), 20, (0, 255, 0), -1)    # 绘制半径为10的实心圆

                for panel in panels:
                    x, y, w, h = panel
                    cv2.circle(frame,(x + w // 2, y + h // 2), 5, (255, 255, 0))  # 绘制灯条中点的小圆

            cv2.imshow("Detected Panels", frame)
            if cv2.waitKey(1) & 0xFF == ord('q'):
                break

        cap.release()
        cv2.destroyAllWindows()
    else:
        print("无效的输入")


if __name__ == "__main__":
    main()